package com.tripcube.ai.service.impl;

import com.tripcube.ai.model.dto.KnowledgeRetrievalDTO;
import com.tripcube.ai.model.dto.RagChatDTO;
import com.tripcube.ai.model.vo.BaseChatVO;
import com.tripcube.ai.model.vo.KnowledgeRetrievalVO;
import com.tripcube.ai.service.KnowledgeRetrievalService;
import com.tripcube.ai.service.RagChatService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import java.util.List;

import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;

/**
 * RAG聊天服务实现类
 * 集成知识库检索功能的聊天服务
 * 
 * @author tripcube
 * @date 2024-01-20
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class RagChatServiceImpl implements RagChatService {
    
    private final ChatClient mysqlMemoryChatClient;
    private final MessageWindowChatMemory messageWindowChatMemory;
    private final KnowledgeRetrievalService knowledgeRetrievalService;
    
    @Override
    public BaseChatVO ragChat(RagChatDTO request) {
        try {
            String enhancedPrompt = buildEnhancedPrompt(request);
            
            String response;
            if (StringUtils.hasText(request.getSystemPrompt())) {
                response = mysqlMemoryChatClient.prompt()
                        .system(request.getSystemPrompt())
                        .user(enhancedPrompt)
                        .advisors(a -> a.param(CONVERSATION_ID, request.getConversationId()))
                        .call().content();
            } else {
                response = mysqlMemoryChatClient.prompt(enhancedPrompt)
                        .advisors(a -> a.param(CONVERSATION_ID, request.getConversationId()))
                        .call().content();
            }
            
            return new BaseChatVO(response, request.getConversationId());
            
        } catch (Exception e) {
            log.error("RAG聊天失败", e);
            return new BaseChatVO("抱歉，处理您的请求时出现了错误，请稍后重试。", request.getConversationId());
        }
    }
    
    @Override
    public Flux<String> ragStreamChat(RagChatDTO request) {
        try {
            String enhancedPrompt = buildEnhancedPrompt(request);
            
            if (StringUtils.hasText(request.getSystemPrompt())) {
                return mysqlMemoryChatClient.prompt()
                        .system(request.getSystemPrompt())
                        .user(enhancedPrompt)
                        .advisors(a -> a.param(CONVERSATION_ID, request.getConversationId()))
                        .stream()
                        .content();
            } else {
                return mysqlMemoryChatClient.prompt(enhancedPrompt)
                        .advisors(a -> a.param(CONVERSATION_ID, request.getConversationId()))
                        .stream()
                        .content();
            }
            
        } catch (Exception e) {
            log.error("RAG流式聊天失败", e);
            return Flux.just("抱歉，处理您的请求时出现了错误，请稍后重试。");
        }
    }
    
    @Override
    public List<Message> getMessages(String conversationId) {
        return messageWindowChatMemory.get(conversationId);
    }
    
    @Override
    public void clearMessages(String conversationId) {
        messageWindowChatMemory.clear(conversationId);
        log.info("已清除会话消息记录，会话ID: {}", conversationId);
    }
    
    @Override
    public BaseChatVO ragChatWithMultipleKnowledgeBases(RagChatDTO request, String[] knowledgeBaseIds) {
        try {
            // 从多个知识库检索相关内容
            KnowledgeRetrievalVO retrievalResult = knowledgeRetrievalService.retrieveFromMultipleKnowledgeBases(
                    request.getQuery(), knowledgeBaseIds, request.getTopK());
            
            String enhancedPrompt = buildEnhancedPromptWithRetrievalResult(request, retrievalResult);
            
            String response;
            if (StringUtils.hasText(request.getSystemPrompt())) {
                response = mysqlMemoryChatClient.prompt()
                        .system(request.getSystemPrompt())
                        .user(enhancedPrompt)
                        .advisors(a -> a.param(CONVERSATION_ID, request.getConversationId()))
                        .call().content();
            } else {
                response = mysqlMemoryChatClient.prompt(enhancedPrompt)
                        .advisors(a -> a.param(CONVERSATION_ID, request.getConversationId()))
                        .call().content();
            }
            
            return new BaseChatVO(response, request.getConversationId());
            
        } catch (Exception e) {
            log.error("多知识库RAG聊天失败", e);
            return new BaseChatVO("抱歉，处理您的请求时出现了错误，请稍后重试。", request.getConversationId());
        }
    }
    
    /**
     * 构建增强的提示词
     */
    private String buildEnhancedPrompt(RagChatDTO request) {
        // 如果未启用知识库检索，直接返回原始查询
        if (!Boolean.TRUE.equals(request.getEnableRetrieval()) ||
            request.getKnowledgeBaseId() == null) {
            return request.getQuery();
        }
        
        try {
            // 构建知识检索请求
            KnowledgeRetrievalDTO retrievalDTO = new KnowledgeRetrievalDTO();
            retrievalDTO.setQuery(request.getQuery());
            retrievalDTO.setKnowledgeBaseId(request.getKnowledgeBaseId());
            retrievalDTO.setTopK(request.getTopK() != null ? request.getTopK() : 5);
            
            // 执行知识检索
            KnowledgeRetrievalVO retrievalResult = knowledgeRetrievalService.retrieveKnowledge(retrievalDTO);
            
            return buildEnhancedPromptWithRetrievalResult(request, retrievalResult);
            
        } catch (Exception e) {
            log.warn("知识检索失败，使用原始查询: {}", e.getMessage());
            return request.getQuery();
        }
    }
    
    /**
     * 基于检索结果构建增强的提示词
     */
    private String buildEnhancedPromptWithRetrievalResult(RagChatDTO request, KnowledgeRetrievalVO retrievalResult) {
        if (retrievalResult == null || retrievalResult.getItems() == null || retrievalResult.getItems().isEmpty()) {
            log.info("未检索到相关知识，使用原始查询");
            return request.getQuery();
        }
        
        StringBuilder contextBuilder = new StringBuilder();
        contextBuilder.append("以下是相关的背景知识，请基于这些信息回答用户的问题：\n\n");
        
        // 添加检索到的知识内容
        for (int i = 0; i < retrievalResult.getItems().size(); i++) {
            KnowledgeRetrievalVO.RetrievalItem item = retrievalResult.getItems().get(i);
            contextBuilder.append(String.format("[知识片段%d] 来源：%s\n", i + 1, item.getDocumentName()));
            contextBuilder.append(item.getContent());
            contextBuilder.append("\n\n");
        }
        
        contextBuilder.append("用户问题：").append(request.getQuery());
        contextBuilder.append("\n\n请基于上述背景知识回答用户问题。如果背景知识中没有相关信息，请明确说明并提供一般性的回答。");
        
        String enhancedPrompt = contextBuilder.toString();
        
        log.info("构建增强提示词完成，检索到{}个相关知识片段", retrievalResult.getItems().size());
        
        return enhancedPrompt;
    }
    
    /**
     * 验证RAG聊天请求参数
     */
    private void validateRagChatRequest(RagChatDTO request) {
        if (!StringUtils.hasText(request.getQuery())) {
            throw new IllegalArgumentException("查询内容不能为空");
        }
        
        if (!StringUtils.hasText(request.getConversationId())) {
            throw new IllegalArgumentException("会话ID不能为空");
        }
        
        if (Boolean.TRUE.equals(request.getEnableRetrieval()) &&
            !StringUtils.hasText(request.getKnowledgeBaseId())) {
            throw new IllegalArgumentException("启用知识库检索时，知识库ID不能为空");
        }
    }
}